package raftpaxos

import (

    "context"
    "math/rand"
    "sync"
    "time"
)

// 定义节点状态常量
const (
    Follower  = iota // 跟随者
    Candidate        // 候选人
    Leader           // 领导者
)

// 定义日志条目结构
type LogEntry struct {
    Term    int    // 任期
    Index   int    // 索引
    Command []byte // 命令
}

// 定义 Raft-Paxos 节点结构
type Node struct {
    id               int              // 节点ID
    state            int              // 节点状态
    currentTerm      int              // 当前任期
    votedFor         int              // 投票给的候选人
    log              []LogEntry       // 节点的日志
    commitIndex      int              // 已提交的日志索引
    lastApplied      int              // 上次已应用的日志索引
    nextIndex        int              // 下一条日志条目索引
    leaderId         int              // 当前领导者的ID
    electionTimeout  time.Duration    // 选举超时时间
    heartbeatTimeout time.Duration    // 心跳超时时间
    mu               sync.Mutex       // 互斥锁
    peers            map[int]*Node    // 节点之间的连接
    ch               chan Event       // 事件通道
    voteCount        int              // 当前选票数
}

// 定义事件类型
type Event struct {
    Type    EventType // 事件类型
    From    int       // 事件发起者ID
    To      int       // 事件接收者ID
    Term    int       // 事件的任期
    Success bool      // 事件是否成功
}

// 定义事件类型常量
type EventType int

const (
    StartElection EventType = iota + 1
    ReceiveRequestVote
    ReceiveAppendEntries
    ReceiveRequestVoteResponse
    ReceiveAppendEntriesResponse
    BecomeLeader
    StepDown
)

// 定义事件结构
type RequestVoteEvent struct {
    CandidateId   int
    CandidateTerm int
    LastLogIndex  int
    LastLogTerm   int
}

type RequestVoteResponseEvent struct {
    VoteGranted   bool
    CurrentTerm   int
    LastLogIndex  int
    LastLogTerm   int
}

type AppendEntriesEvent struct {
    LeaderId     int
    PrevLogIndex int
    PrevLogTerm  int
    Entries      []LogEntry
    LeaderCommit int
}

type AppendEntriesResponseEvent struct {
    Success      bool
    CurrentTerm  int
    LastLogIndex int
}

// Node 构造器函数
func NewNode(id int, peers []int, eventChannel chan Event) *Node {
    // 初始化节点
    node := &Node{
        id:               id,
        state:            Follower,
        peers:            make(map[int]*Node),
        ch:               eventChannel,
        electionTimeout:  time.Duration(rand.Intn(300) + 150) * time.Millisecond,
        heartbeatTimeout: time.Duration(150) * time.Millisecond,
    }

    // 建立与其他节点的连接
    for _, peer := range peers {
        if peer == id {
            continue
        }
        node.peers[peer] = nil // 注意：在实际实现中，应使用网络或通道与其他节点通信
    }

    return node
}

type RaftPaxos struct {
    nodes []*Node
}

func NewRaftPaxos(nodeIds []int) (*RaftPaxos, error) {
    eventChannel := make(chan Event, 100)

    raftPaxos := &RaftPaxos{
        nodes: make([]*Node, len(nodeIds)),
    }

    var wg sync.WaitGroup
    errorChannel := make(chan error, len(nodeIds)) // 用于收集错误

    for i, nodeId := range nodeIds {
        wg.Add(1)
        node := NewNode(nodeId, nodeIds, eventChannel)
        raftPaxos.nodes[i] = node

        // 修改为使用上下文管理goroutine
        ctx, cancel := context.WithCancel(context.Background())
        go func(ctx context.Context, node *Node) {
            defer wg.Done()
            if err := node.run(ctx); err != nil {
                errorChannel <- err
            }
        }(ctx, node)

        // 可以取消node的运行（例如在系统关闭时）
        go func() {
            wg.Wait()
            cancel()
        }()
    }

    // 等待所有节点启动完成，或直到发生错误
    select {
    case err := <-errorChannel:
        close(errorChannel)
        wg.Wait()
        return nil, err // 返回第一个遇到的错误
    case <-ctx.Done():
        // 可能需要处理上下文被取消的情况
    }

    return raftPaxos, nil
}


// Node 运行函数
func (n *Node) run() {
    for {
        // 监听事件通道
        event := <-n.ch
        // 根据事件类型进行处理
        switch e := event.(type) {
        case RequestVoteEvent:
            n.handleRequestVoteEvent(e)
        case AppendEntriesEvent:
            n.handleAppendEntriesEvent(e)
        case RequestVoteResponseEvent:
            n.handleRequestVoteResponseEvent(e)
        case AppendEntriesResponseEvent:
            n.handleAppendEntriesResponseEvent(e)
        }
    }
}

// 处理 RequestVoteEvent 事件
func (n *Node) handleRequestVoteEvent(event RequestVoteEvent) {
    n.mu.Lock()
    defer n.mu.Unlock()

    // 如果收到的任期更高，更新状态
    if event.CandidateTerm > n.currentTerm {
        n.currentTerm = event.CandidateTerm
        n.votedFor = 0
        n.state = Follower
    }

    granted := false
    // 根据当前任期和投票情况决定是否投票给候选人
    if event.CandidateTerm == n.currentTerm && (n.votedFor == 0 || n.votedFor == event.CandidateId) {
        n.votedFor = event.CandidateId
        granted = true
    }

    // 发送投票结果响应
    response := RequestVoteResponseEvent{
        VoteGranted: granted,
        CurrentTerm: n.currentTerm,
        LastLogIndex: len(n.log) - 1,
        LastLogTerm: n.log[len(n.log) - 1].Term,
    }
    n.peers[event.CandidateId].ch <- response
}

// 处理 AppendEntriesEvent 事件
func (n *Node) handleAppendEntriesEvent(event AppendEntriesEvent) {
    n.mu.Lock()
    defer n.mu.Unlock()

    // 如果收到的任期更高，更新状态
    if event.LeaderId != n.id {
        n.leaderId = event.LeaderId
        n.state = Follower
    }

    // 判断日志同步状态
    success := false
    if event.PrevLogIndex == len(n.log) - 1 && event.PrevLogTerm == n.log[len(n.log) - 1].Term {
        n.log = append(n.log, event.Entries...)
        n.commitIndex = event.LeaderCommit
        success = true
    }

    // 发送 AppendEntriesResponse 事件
    response := AppendEntriesResponseEvent{
        Success: success,
        CurrentTerm: n.currentTerm,
        LastLogIndex: len(n.log) - 1,
    }
    n.peers[event.LeaderId].ch <- response
}

// 处理 RequestVoteResponseEvent 事件
func (n *Node) handleRequestVoteResponseEvent(event RequestVoteResponseEvent) {
    n.mu.Lock()
    defer n.mu.Unlock()

    // 如果当前任期与响应的任期匹配，并且响应投票成功
    if event.CurrentTerm == n.currentTerm && event.VoteGranted {
        n.voteCount++
        // 如果获得超过半数选票，成为领导者
        if n.voteCount > len(n.peers)/2 {
            n.becomeLeader()
        }
    }
}

// 成为领导者
func (n *Node) becomeLeader() {
    n.mu.Lock()
    defer n.mu.Unlock()

    n.state = Leader
    n.leaderId = n.id
    n.nextIndex = len(n.log)
    n.commitIndex = len(n.log) - 1

    // 开始定期发送心跳消息
    go n.sendHeartbeats()
}

// 发送心跳
func (n *Node) sendHeartbeats() {
    for {
        n.sendAppendEntries()
        time.Sleep(n.heartbeatTimeout)
    }
}

// 发送 AppendEntries 事件
func (n *Node) sendAppendEntries() {
    n.mu.Lock()
    defer n.mu.Unlock()

    // 构造 AppendEntries 事件并发送给其他节点
    event := AppendEntriesEvent{
        LeaderId: n.id,
        PrevLogIndex: len(n.log) - 1,
        PrevLogTerm: n.log[len(n.log) - 1].Term,
        Entries: []LogEntry{}, // 可以根据实际需求传递日志条目
        LeaderCommit: n.commitIndex,
    }

    // 发送事件给其他节点
    for _, peer := range n.peers {
        if peer == n.id {
            continue
        }
        peer.ch <- event
    }
}

// 处理 AppendEntriesResponseEvent 事件
func (n *Node) handleAppendEntriesResponseEvent(event AppendEntriesResponseEvent) {
    // 处理追加日志响应事件
    // 如果追加成功，更新 nextIndex
    if event.Success {
        n.nextIndex = event.LastLogIndex + 1
    } else {
        // 如果追加失败，尝试回退日志同步
        n.nextIndex--
        n.sendAppendEntries()
    }
}
